/**
 * Copyright 2004 - 2007 Blue Bamboo International Inc. 
 *           All rights reserved.
 *
 *
 * Blue Bamboo PROPRIETARY/CONFIDENTIAL.
 *
 */

package com.bambamboo.st.socket.server.netty;

import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.security.GeneralSecurityException;
import java.security.KeyStore;

import javax.net.ssl.KeyManager;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509TrustManager;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.bambamboo.st.util.StringUtil;

/**
 * DOCME
 * 
 * @author tanxuqing
 * @date Dec 24, 2009
 * @since
 */
public class SSLContextProvider {

	private static final Logger logger = LoggerFactory.getLogger(SSLContextProvider.class);

	private String trustStoreFileName = null;
	private String trustStorePassword = null;
	private String keyStoreFileName = null;
	private String keyStorePassword = null;
	private boolean disableCertValid = false;

	public SSLContextProvider(String trustStoreFileName, String trustStorePassword, String keyStoreFileName,
			String keyStorePassword, boolean disableCertValid) {
		this.trustStoreFileName = trustStoreFileName;
		this.trustStorePassword = trustStorePassword;
		this.keyStoreFileName = keyStoreFileName;
		this.keyStorePassword = keyStorePassword;
		this.disableCertValid = disableCertValid;
	}

	public SSLContext getSSLContext() throws IOException, GeneralSecurityException {
		// Call getTrustManagers to get suitable trust managers
		TrustManager[] tms = getTrustManagers();

		// Call getKeyManagers (from CustomKeyStoreClient) to get suitable
		// key managers
		KeyManager[] kms = getKeyManagers();

		// Next construct and initialize a SSLContext with the KeyStore and
		// the TrustStore. We use the default SecureRandom.
		SSLContext context = SSLContext.getInstance("SSL");
		context.init(kms, tms, null);

		return context;
	}

	protected KeyManager[] getKeyManagers() throws GeneralSecurityException {
		KeyManager[] kms = null;
		if (StringUtil.isEmpty(keyStoreFileName) || StringUtil.isEmpty(keyStorePassword)) {
			return null;
		}
		// First, get the default KeyManagerFactory.
		try {
			String alg = KeyManagerFactory.getDefaultAlgorithm();
			KeyManagerFactory kmFact = KeyManagerFactory.getInstance(alg);

			// Next, set up the KeyStore to use. We need to load the file into
			// a KeyStore instance.
			FileInputStream fis = new FileInputStream(keyStoreFileName);
			KeyStore ks = KeyStore.getInstance("JKS");
			ks.load(fis, keyStorePassword.toCharArray());
			fis.close();

			// Now we initialize the TrustManagerFactory with this KeyStore
			kmFact.init(ks, keyStorePassword.toCharArray());

			// And now get the TrustManagers
			kms = kmFact.getKeyManagers();
		}
		catch (FileNotFoundException fe) {
			logger.error("Invalid KeyStore File Name: " + keyStoreFileName, fe);
			kms = null;
		}
		catch (Exception e) {
			logger.error("IOExceptin of KeyStore File Name: " + keyStoreFileName, e);
			kms = null;
		}
		return kms;
	}

	protected TrustManager[] getTrustManagers() throws GeneralSecurityException {
		TrustManager[] tms = null;
		if (disableCertValid) {
			// trust all the SSL server
			return new TrustManager[] {
				new X509TrustManager() {

					public java.security.cert.X509Certificate[] getAcceptedIssuers() {
						return null;
					}

					public void checkClientTrusted(java.security.cert.X509Certificate[] certs, String authType) {
						// No need to implement.
					}

					public void checkServerTrusted(java.security.cert.X509Certificate[] certs, String authType) {
						// No need to implement.
					}
				}
			};
		}
		// update by Hans, by default trust all the SSL server
		if (StringUtil.isEmpty(trustStoreFileName) || StringUtil.isEmpty(trustStorePassword)) {
			return null;// use the system default truct managers
		}

		try {
			// First, get the default TrustManagerFactory.
			String alg = TrustManagerFactory.getDefaultAlgorithm();
			TrustManagerFactory tmFact = TrustManagerFactory.getInstance(alg);

			// Next, set up the TrustStore to use. We need to load the file into
			// a KeyStore instance.
			FileInputStream fis = new FileInputStream(trustStoreFileName);
			KeyStore ks = KeyStore.getInstance("JKS");
			ks.load(fis, trustStorePassword.toCharArray());
			fis.close();

			// Now we initialize the TrustManagerFactory with this KeyStore
			tmFact.init(ks);

			// And now get the TrustManagers
			tms = tmFact.getTrustManagers();
		}
		catch (FileNotFoundException fe) {
			logger.error("Invalid TrustStore File Name: " + trustStoreFileName, fe);
			tms = null;
		}
		catch (Exception e) {
			logger.error("IOExceptin of TrustStore File Name: " + trustStoreFileName, e);
			tms = null;
		}
		return tms;
	}
}
